using System;
using System.Runtime.InteropServices;
using kcp;

// 引入 mimalloc 的 C 函数
internal static class Mimalloc {
    [DllImport("mimalloc", CallingConvention = CallingConvention.Cdecl)]
    public static extern IntPtr mi_malloc(ulong size);

    [DllImport("mimalloc", CallingConvention = CallingConvention.Cdecl)]
    public static extern void mi_free(IntPtr ptr);
}

public static class KcpMimallocSetup {
    // 自定义分配函数：调用 mimalloc 的 mi_malloc
    private static unsafe void* MimallocMalloc(nuint size) {
        IntPtr ptr = Mimalloc.mi_malloc((ulong)size);
        return (void*)ptr;
    }

    // 自定义释放函数：调用 mimalloc 的 mi_free
    private static unsafe void MimallocFree(void* ptr) {
        Mimalloc.mi_free((IntPtr)ptr);
    }
    
    // 初始化 KCP 并设置 mimalloc 作为分配器
    public static unsafe KcpTransport.LowLevel.IKCPCB* CreateKcpWithMimalloc(uint conv, void* user) {
        // 设置自定义分配器
        delegate* managed<nuint, void*> mallocDelegate = &MimallocMalloc;
        delegate* managed<void*, void> freeDelegate = &MimallocFree;
        KCP.ikcp_allocator(mallocDelegate, freeDelegate);

        // 创建 KCP 实例（此时会使用 mimalloc 分配内存）
        return KcpTransport.LowLevel.KcpMethods.ikcp_create(conv, user);
    }

    // 初始化 KCP 并设置 mimalloc 作为分配器
    public static unsafe IKCPCB* CreateKcpWithMimalloc2(uint conv, void* user) {
        // 设置自定义分配器
        delegate* managed<nuint, void*> mallocDelegate = &MimallocMalloc;
        delegate* managed<void*, void> freeDelegate = &MimallocFree;
        KCP.ikcp_allocator(mallocDelegate, freeDelegate);

        // 创建 KCP 实例（此时会使用 mimalloc 分配内存）
        return KCP.ikcp_create(conv, user);
    }
}